import gym
from gym.envs.registration import register
import os
from ril_env.atari_wrapper import wrap_deepmind

os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = 'hide'  # 隐藏 pygame 版本信息

register(
    id='FlappyBird-v1',
    entry_point='ril_env.envs:FlappyBirdEnv',
    reward_threshold=100
)

register(
    id='Breakout-v1',
    entry_point='ril_env.envs:BreakoutEnv',
    reward_threshold=400
)

register(
    id='SpaceWar-v0',
    entry_point='ril_env.envs:SpaceWarEnv',
    reward_threshold=500
)

register(
    id='GridWorld-v0',
    entry_point='ril_env.envs:GridWorldEnvV0',
    reward_threshold=500
)

register(
    id='GridWorld-v1',
    entry_point='ril_env.envs:GridWorldEnvV1',
    reward_threshold=500
)

"""
param for wrap_deepmind:
(episode_life=True, clip_rewards=True,
frame_stack=4, scale=False, warp_frame=True,
noop_reset=True, max_and_skip=4, contain_origin=False)
"""
wrap_envs = {
    'BreakoutNoFrameskip-v4': {
        'train': [],
        'watch': []
    },
    'PongNoFrameskip-v4': {
        'train': [],
        'watch': []
    },
    'FlappyBird-v1': {
        'train': [False, False, 4, False, True, True, 4, True],
        'watch': [False, False, 4, False, True, True, 4, True]
    },
    'SpaceWar-v0': {
        'train': [False, False, 4, False, True, True, 1, True],
        'watch': [False, False, 4, False, True, True, 1, True]
    },
    'GridWorld-v0': {
        'train': [False, False, 4, False, False, False, None],
        'watch': [False, False, 4, False, False, False, None]
    }
}


def env_factory(task: str, watch: bool = False, ril: bool = False):
    if task in wrap_envs:
        param = wrap_envs[task]['watch' if watch else 'train']
        if not ril and len(param) > 7 and param[7]:
            param[7] = False
        return wrap_deepmind(task, *param)
    return gym.make(task)


def make_atari_env(args):
    return env_factory(args.task, ril=args.ril)


def make_atari_env_watch(args):
    return env_factory(args.task, watch=True, ril=args.ril)
